package org.example.myleet.p698;

import java.util.HashSet;
import java.util.Set;

public class Solution {

    private Set<Integer> existSelection;

    public boolean canPartitionKSubsets(int[] nums, int k) {
        int totalSelection = 1;
        for (int i = 1; i < nums.length; ++i) {
            totalSelection <<= 1;
            totalSelection += 1;
        }
        existSelection = new HashSet<>(totalSelection);
        int sum = 0, max = Integer.MIN_VALUE;
        for (int num : nums) {
            sum += num;
            max = Math.max(max, num);
        }
        if (sum % k != 0) {
            return false;
        }
        int average = sum / k;
        if (max > average) {
            return false;
        }
        return backtrack(nums, average, totalSelection, 0, 0);
    }

    private boolean backtrack(int[] nums, int average, int end, int selection, int setSum) {
        if (setSum == 0 && end == selection) {
            return true;
        }
        for (int i = 0; i < nums.length; ++i) {
            int select = 1 << i;
            if ((selection & select) != 0) {
                continue;
            }
            int newSelection = selection | select;
            if (existSelection.contains(newSelection)) {
                continue;
            }
            existSelection.add(newSelection);
            int newSetSum = setSum + nums[i];
            if (newSetSum > average) {
                continue;
            }
            if (newSetSum == average) {
                newSetSum = 0;
            }
            if (backtrack(nums, average, end, newSelection, newSetSum)) {
                return true;
            }
        }
        return false;
    }
}
